%% osl_africa.m
%
% AfRICA - ArteFact Rejection using Independent Component Analysis
%
% Syntax: fname_out=osl_africa(S)
% S needs to contain:
%   -  fname: full name of input SPM object (inc. path if not in present
%      working directory) e.g. S.fname = '/home/data/spm8_mydata'.
% OPTIONAL inputs:
%   -  just_ica: If set, AfRICA will skip the classification stage and only
%      run the ICA decomposition, save that to disc but not attempt to
%      clean up the SPM object. DEFAULT = 0;
%   -  do_plots: set to 1 to output summary plots of artefact components.
%      set to 0 to switch off plotting.DEFAULT = 1.
%   -  ica_params: a matlab structure containing the following fields:
%           - num_ics: Number of indpendent components sought: DEFAULT =
%           150.
%           - last_eig: order of PCA prewhitening: DEFAULT = num_ics.
%           - nonlinearity: DEFAULT = 'tanh'
%           consult fastica documentation for more information on ica_params.
%    - used_maxfilter: set to 1 if using Maxfiltered data. DEFAULT = 0;
%    - ica_file: user specified file name for the the ICA decomposition. If
%      ica_file does not exist then AfRICA will save the ICA decomposition
%      here. If it does then AfRICA will attempt to read the existing
%      result in.
%    - ident.func: function handle for user-specified function for selecting
%      bad components. Default @IDENTIFY_ARTEFACTUAL_COMPONENTS_MANUAL. Use
%      this as a template ofr custom functions.
%    - ident: contains all settings for ident.func
% outputs
%   - fname_out: the name of the output SPM object. The prefix 'A' is
%     attached to any objects that have been cleaned up with AfRICA. 
%     e.g. fname_out = '/home/data/Aspm8_mydata.dat'
%
%     Example:
%     S=[];
%     S.fname='spm8_mydata';                     
%     S.logfile = 1;                                                                 
%     S.ica_file = [S.fname '_preproc_ica_results']; 
%     S.used_maxfilter = 1;                         
%     
%     S.eyetracker_file='my_eyetracker_file.txt';    
%     S.eyetracker_chan=319;                          
%     S.do_plots=1;                                   
%     S.manual_approval=1;                                                             
%                                                                                             
%     S.ident=ident;
%     S.ident.func=@func;
%    
%     [spm_files_new]=osl_africa(S);
%
% AfRICA ignores bad trials defined by D.reject and bad periods defined by
% 'BadEpoch' events using OSLVIEW. Do not rest these periods to "good"
% after AfRICA.
%
% HL+AB 131112
% henry.luckhoo@trinity.ox.ac.uk

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% MAIN FUNCTION %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

function [fname_out, fig_handles, fig_names, fig_titles, S]=osl_africa(S)

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% SETUP %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

try
    fname=S.fname;
    D=spm_eeg_load(S.fname);
    S.D=D;
catch
    error('Can not find SPM file from S.fname')
end

if not(isfield(S,'modality'))   % added by DM
    S.modality='MEG';
end

[dir,nam,ext]=fileparts(fname);
fname_out=[dir '/A' nam '.dat'];

if ~isfield(S,'to_do'); S.S.to_do = [1 1 1]; end

if isfield(S,'logfile') && S.logfile == 1
  logdir = [dir '/Africa_logs/'];
  if ~isdir(logdir), mkdir(logdir); end
  if exist([logdir nam '_log.txt'],'file')
    unix(['rm ' logdir nam '_log.txt']);
  end
  	diary([logdir nam '_log.txt']);
end

if ~isfield(S,'used_maxfilter'); S.used_maxfilter = 0;end

% Allow user-specified functions to be used to identify bad components.
if isfield(S.ident,'func') && isa(S.ident.func,'function_handle'); 
    identfun = S.ident.func; 
    msg = sprintf('\n%s%s%s\n%','User-specified function for identifying bad components "', func2str(identfun) ,'.m" being used.');
    fprintf(msg);
else
    identfun = @identify_artefactual_components_auto;
    msg = sprintf('\n%s%s%s\n%','Standard OSL sub-function for identifying bad components "', func2str(identfun) ,'" being used.');
    fprintf(msg)
end

fig_handles=[]; 
fig_names=[];
fig_titles=[];

%%%%%%%%%%%%%%%%%%%%%%%%% PERFORM SENSOR SPACE ICA %%%%%%%%%%%%%%%%%%%%%%%%

if S.to_do(1)
    S.ica_res = perform_sensorspace_ica(S);
    if isfield(S,'ica_file')
        save(S.ica_file,'S');
        msg = sprintf('\n%s%s\n%','Saving ICA results to ', [S.ica_file '.mat.']);
        fprintf(msg);
    else
        msg = sprintf('\n%s\n%','Results not being saved.');
        fprintf(msg);
    end
end
%%%%%%%%%%%%%%%%%%%%%%%%%% CLASSIFY BAD COMPONENTS %%%%%%%%%%%%%%%%%%%%%%%%

if S.to_do(2)
    if isfield(S,'ica_file')
        if (exist([S.ica_file '.mat'],'file')==2)
            Snew = load(S.ica_file);
            if isfield(S,'ident'); Snew.S.ident=S.ident; end;
            if isfield(S,'do_plots');        Snew.S.do_plots        = S.do_plots;        end
            if isfield(S,'to_do');           Snew.S.to_do           = S.to_do;           end
            S = Snew.S;
            
            msg = sprintf('\n%s%s\n%','Loading previous ICA results from ', [S.ica_file '.mat.']);
            fprintf(msg);
            if(S.do_plots)
                [S.ica_res.bad_components, fig_handles, fig_names, fig_titles] = feval(identfun,S);
            else
                S.ica_res.bad_components = feval(identfun,S);
            end;   
            save(S.ica_file,'S');
            msg = sprintf('\n%s%s\n%','Saving bad component selection to ', [S.ica_file '.mat.']);
            fprintf(msg);
        else
            if isfield(S,'ica_res')
                if(S.do_plots)
                    [S.ica_res.bad_components, fig_handles, fig_names, fig_titles] = feval(identfun,S);
                else
                    S.ica_res.bad_components = feval(identfun,S);
                end;
                ica_res=S.ica_res;
                save(S.ica_file,'ica_res');
                msg = sprintf('\n%s%s\n%','Saving bad component selection to ', [S.ica_file '.mat']);
                fprintf(msg);
            else
                error('ICA decomposition needs to be run. Set S.to_do(1) = 1 and try again');
            end
        end
    else
        
        if(S.do_plots)
            [S.ica_res.bad_components, fig_handles, fig_names, fig_titles] = feval(identfun,S);
        else
            S.ica_res.bad_components = feval(identfun,S);
        end;
    end
end

%%%%%%%%%%%%%%%%%% REMOVE BAD COMPONENTS FROM THE DATA %%%%%%%%%%%%%%%%%%%%
if S.to_do(3)
    if ~isfield(S,'ica_res')
        try
            load(S.ica_file);
        catch
            error('Unable to load ICA decomposition from disk. Try rerunning all AfRICA stages');
        end
    end
    fname_out = remove_bad_components(S);
else
    fname_out = [];
end

if isfield(S,'logfile') && S.logfile == 1,
    diary off;
end

end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% SUB-FUNCTIONS %%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% perform_sensorspace_ica - subfunction to decompose meg data using fastICA

function res = perform_sensorspace_ica(S)

%%%%%%%%%%%%%%%%%%%%%%%% LOAD AND PREPARE MEG DATA %%%%%%%%%%%%%%%%%%%%%%%%

D       = spm_eeg_load(S.fname);


if strcmp(S.modality,'EEG')
    chan_inds=setdiff(find(any([strcmp(D.chantype,'EEG')],1)),D.badchannels);
else
% 1.) Select good MEG Channels
chan_inds=setdiff(find(any([strcmp(D.chantype,'MEGMAG');strcmp(D.chantype,'MEGPLANAR');strcmp(D.chantype,'MEGGRAD')],1)),D.badchannels);
end
meg_dat=D(chan_inds,:,:);

% 2.) Remove bad trials/any trial structure
meg_dat=meg_dat(:,:,setdiff(1:D.ntrials, find(D.reject)));
c=1; %Ntrials=size(meg_dat,3);
meg_dat_rs=zeros([size(meg_dat,1), size(meg_dat,2)*size(meg_dat,3)]);
for i=1:size(meg_dat,3)
    meg_dat_rs(:,c:c+D.nsamples-1)=meg_dat(:,:,i);
    c=c+D.nsamples;
end
meg_dat=meg_dat_rs;
clear meg_dat_rs;

% 3.) Remove bad segments
if D.ntrials==1;
    t = D.time;
    badsections = false(1,D.nsamples);
    Events = D.events;
    if ~isempty(Events)
        Events = Events(strcmp({Events.type},'BadEpoch'));
        for ev = 1:numel(Events)
            badsections = badsections | t >= Events(ev).time & t < (Events(ev).time+Events(ev).duration);
        end
    end
    meg_dat(:,badsections)=[];
end

%%%%%%%%%%%%%%%%%%%% APPLY MAXFILTER SPECIFIC SETTINGS %%%%%%%%%%%%%%%%%%%%

if isfield(S,'used_maxfilter') && S.used_maxfilter
    num_ics_default=62;
    mag_cutoff  = 62;
    plan_cutoff = 62;
else
    num_ics_default=150;
    mag_cutoff  = sum(strcmp(D.chantype(chan_inds),'MEGMAG'))    - 5;
    plan_cutoff = sum(strcmp(D.chantype(chan_inds),'MEGPLANAR')) - 5;
end

%%%%%%%%%%%%%%%%%%%%%%%%% SET FASTICA PARAMETERS %%%%%%%%%%%%%%%%%%%%%%%%%%

if isfield(S,'ica_params')
    if isfield(S.ica_params,'num_ics');num_ics=S.ica_params.num_ics;else num_ics=num_ics_default; end
    if isfield(S.ica_params,'last_eig');last_eig=S.ica_params.last_eig;else last_eig=num_ics_default; end
    if isfield(S.ica_params,'nonlinearity');nonlinearity=S.ica_params.nonlinearity;else nonlinearity='tanh'; end
    if isfield(S.ica_params,'approach');ica_approach=S.ica_params.approach;else ica_approach='symm'; end  % added by DM
    if isfield(S.ica_params,'max_iterations');max_iter=S.ica_params.max_iterations;else max_iter=1000; end
else
    num_ics=num_ics_default;
    last_eig=num_ics_default;
    nonlinearity='tanh';
    ica_approach='symm';
    max_iter=1000;
end

num_ics=min(num_ics,size(meg_dat,1));  % added by DM
last_eig=min(last_eig,size(meg_dat,1));  % added by DM

%%%%%%%%%%%%%%%%%%%%  MINIMUM EIGENVALUE NORMALISATION %%%%%%%%%%%%%%%%%%%%

if strcmp(S.modality,'EEG')  % added by DM
    norm_vec=max(abs(meg_dat(:)))/1000*ones(size(meg_dat,1),1);
    
else
    
    norm_vec=ones(numel(chan_inds),1);
    if any(strcmp(D.chantype,'MEGMAG')) && any(strcmp(D.chantype,'MEGPLANAR'))
        mag_min_eig=svd(cov(meg_dat(strcmp(D.chantype(chan_inds),'MEGMAG'),:)')); mag_min_eig=mean(mag_min_eig(mag_cutoff-2:mag_cutoff));
        plan_min_eig=svd(cov(meg_dat(strcmp(D.chantype(chan_inds),'MEGPLANAR'),:)')); plan_min_eig=mean(plan_min_eig(plan_cutoff-2:plan_cutoff));
        norm_vec(strcmp(D.chantype(chan_inds),'MEGMAG'))=mag_min_eig; norm_vec(strcmp(D.chantype(chan_inds),'MEGPLANAR'))=plan_min_eig;
    else
        norm_vec=norm_vec*min(svd(cov(meg_dat(strcmp(D.chantype(chan_inds),'MEGGRAD'),:)')));
    end
    norm_vec=sqrt(norm_vec);
    
end

dat_in=meg_dat./repmat(norm_vec,1,size(meg_dat,2));

if S.do_plots
    figure; semilogy(svd(cov(meg_dat')));
    ho;semilogy(svd(cov(dat_in')),'r--');
    title('Raw and normalised eigen spectra'); legend('Raw', 'Normalised');
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%% ICA DECOMPOSITION %%%%%%%%%%%%%%%%%%%%%%%%%%%%

[tc,sm,~]=fastica(dat_in,'g',nonlinearity,'lastEig',last_eig, 'numOfIC', num_ics,'approach',ica_approach, 'maxNumIterations',max_iter); % changed by DM

if num_ics ~= size(tc,1)
    msg = sprintf('\n%s%d%s%d%s\n%','Data dimensionality insufficient to support ', num_ics, ' components. Number of components has been reduced to ', size(tc,1), '.');
    fprintf(msg);
    num_ics = size(tc,1);
end

ica_res=[];
ica_res.ica_params.num_ics = num_ics;
ica_res.ica_params.last_eig = last_eig;
ica_res.ica_params.nonlinearity = nonlinearity;
ica_res.D=S.fname;
ica_res.tc=tc;
ica_res.sm=sm.*repmat(norm_vec,1,num_ics);

%%%%%%%%%%%%%%%%%%%%%%%%% ESTIMATE MISSING CHANNELS %%%%%%%%%%%%%%%%%%%%%%% 

if strcmp(S.modality,'EEG')  % added by DM
sm_full = zeros(numel(find(any([strcmp(D.chantype,'EEG')],1))),ica_res.ica_params.num_ics);
map_inds(find(any([strcmp(D.chantype,'EEG')],1))) = 1:numel(find(any([strcmp(D.chantype,'EEG')],1)));
else
sm_full = zeros(numel(find(any([strcmp(D.chantype,'MEGMAG');strcmp(D.chantype,'MEGPLANAR');strcmp(D.chantype,'MEGGRAD')],1))),ica_res.ica_params.num_ics);
map_inds(find(any([strcmp(D.chantype,'MEGMAG');strcmp(D.chantype,'MEGPLANAR');strcmp(D.chantype,'MEGGRAD')],1))) = 1:numel(find(any([strcmp(D.chantype,'MEGMAG');strcmp(D.chantype,'MEGPLANAR');strcmp(D.chantype,'MEGGRAD')],1)));
end
sm_full(map_inds(chan_inds),:)=ica_res.sm;

excluded_data=D(find(D.badchannels),:,setdiff(1:D.ntrials, find(D.reject)));
c=1;
excluded_data_rs=zeros([size(excluded_data,1), size(excluded_data,2)*size(excluded_data,3)]);
for i=1:size(excluded_data,3)
    excluded_data_rs(:,c:c+D.nsamples-1)=excluded_data(:,:,i);
    c=c+D.nsamples;
end
excluded_data=excluded_data_rs;
clear excluded_data_rs;
if D.ntrials==1;
    t = D.time;
    badsections = false(1,D.nsamples);
    Events = D.events;
    if ~isempty(Events)
        Events = Events(strcmp({Events.type},'BadEpoch'));
        for ev = 1:numel(Events)
            badsections = badsections | t >= Events(ev).time & t < (Events(ev).time+Events(ev).duration);
        end
    end
    excluded_data(:,badsections)=[];
end

sm_full(D.badchannels,:) = excluded_data*pinv(ica_res.tc);
ica_res.sm = sm_full;

%%%%%%%%%%%%%%%%%%%%%%%%%% ESTIMATE MISSING EPOCHS %%%%%%%%%%%%%%%%%%%%%%%% 

excluded_timepoints = false(1,D.ntrials*D.nsamples);
for i=1:D.ntrials
    if D.reject(i)
        excluded_timepoints(1+(i-1)*D.nsamples:i*D.nsamples) = 1;
    end
end
excluded_timepoints(badsections)=true;

if strcmp(S.modality,'EEG')  % added by DM
    excluded_data = D(find(any([strcmp(D.chantype,'EEG')],1)),:,:);
else
    excluded_data = D(find(any([strcmp(D.chantype,'MEGMAG');strcmp(D.chantype,'MEGPLANAR');strcmp(D.chantype,'MEGGRAD')],1)),:,:);
end

c=1;excluded_data_rs=zeros([size(excluded_data,1), size(excluded_data,2)*size(excluded_data,3)]);
for i=1:size(excluded_data,3)
    excluded_data_rs(:,c:c+D.nsamples-1)=excluded_data(:,:,i);
    c=c+D.nsamples;
end
excluded_data=excluded_data_rs;
clear excluded_data_rs;
excluded_data = excluded_data(:,excluded_timepoints);

tc_full = zeros(ica_res.ica_params.num_ics,D.ntrials*D.nsamples);
tc_full(:,~excluded_timepoints) = ica_res.tc;
tc_full(:,excluded_timepoints) = (excluded_data'*pinv(sm_full'))';
ica_res.tc = tc_full;


res = ica_res;

end

%% remove_bad_components - function to subtract the bad components from the
%% MEG data via spm_eeg_montage.m

function res = remove_bad_components(S)

%%%%%%%%%%%%%%%%%%%%%%%% LOAD AND PREPARE MEG DATA %%%%%%%%%%%%%%%%%%%%%%%%

D       = spm_eeg_load(S.fname);

if strcmp(S.modality,'EEG')  % added by DM
    chan_inds=setdiff(find(any([strcmp(D.chantype,'EEG')],1)),D.badchannels);
else
    chan_inds=setdiff(find(any([strcmp(D.chantype,'MEGMAG');strcmp(D.chantype,'MEGPLANAR');strcmp(D.chantype,'MEGGRAD')],1)),D.badchannels);
end
badind = D.badchannels;
bad_components = S.ica_res.bad_components;
meg_dat=D(chan_inds,:,:);

c=1;
meg_dat_rs=zeros([size(meg_dat,1), size(meg_dat,2)*size(meg_dat,3)]);
for i=1:size(meg_dat,3)
    meg_dat_rs(:,c:c+D.nsamples-1)=meg_dat(:,:,i);
    c=c+D.nsamples;
end
meg_dat=meg_dat_rs;
clear meg_dat_rs;

%%%%%%%%%%%%%%%%% SAVE A COPY OF THE TRA MATRIX IF NEEDED %%%%%%%%%%%%%%%%%

if strcmp(S.modality,'MEG')   % added by DM
D = save_raw_tra_to_D(D);
end

%%%%%%%%%%%%%%%%%%% REMOVE BAD COMPONENTS USING MONTAGE %%%%%%%%%%%%%%%%%%%

sm = S.ica_res.sm(chan_inds,:);
tc = S.ica_res.tc;

if strcmp(S.modality,'EEG')   % added by DM
    use_montage=0;
else
    use_montage = 1;
end

if ~isempty(bad_components)
    if use_montage
        dat_inv = pinv_plus(meg_dat',S.ica_res.ica_params.num_ics);
        tra = (eye(numel(chan_inds)) - dat_inv*(tc(bad_components,:)'*sm(:,bad_components)'))';
        
        montage           =  [];
        montage.tra       =  tra;
        montage.labelnew  =  D.chanlabels(chan_inds);
        montage.labelorg  =  D.chanlabels(chan_inds);
        
        if ~isempty(badind)
            montage.labelorg = [montage.labelorg(:); D.chanlabels(badind)']; 
            montage.tra(end, end+length(badind)) = 0;
        end
        
        S_montage                =  [];
        S_montage.D              =  S.fname;
        S_montage.montage        =  montage;
        S_montage.keepothers     =  'yes';
        S_montage.updatehistory  =  1;
        
        Dnew = spm_eeg_montage(S_montage);
        
        %%%%%%%%%%%%%%%%%%%%%%% SAVE NEW SPM OBJECT TO DISK %%%%%%%%%%%%%%%%%%%%%%%
        
        S_copy         = [];
        S_copy.D       = Dnew;
        S_copy.newname = fullfile(D.path, ['A' D.fname]);
        Dnew2=spm_eeg_copy(S_copy);
        Dnew.delete;
    else
        [dir,nam,~]=fileparts(fullfile(D.path,D.fname));
        fname_out=[dir '/A' nam '.dat'];
        meg_dat_clean=meg_dat-(sm(:,bad_components)*tc(bad_components,:));
        
        Dnew2=clone(D,fname_out,size(D));
        Dnew2(chan_inds,:)=meg_dat_clean;   % changed by DM
%         Dnew2(1:306,:)=meg_dat_clean;
%         Dnew2(307:D.nchannels,:)=D(307:D.nchannels,:,:);
        Dnew2.save;
    end
else
    Dnew2=D;
    msg = sprintf('\n%s\n%s\n%','No bad components have been selected. No de-noising has been applied to ', fullfile(D.path, D.fname));
    fprintf(msg);
end

res = fullfile(Dnew2.path, Dnew2.fname);

end